package IC.Context;

import IC.AST.*;
import IC.AST.PrimitiveType;
import IC.BinaryOps;
import IC.Symbols.Kind;
import IC.Symbols.Symbol;
import IC.Symbols.SymbolTable;
import IC.Symbols.SymbolTableKind;
import IC.Types.*;

import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Created by IntelliJ IDEA.
 * User: nimrod
 * Date: 10/01/12
 * Time: 14:23
 * To change this template use File | Settings | File Templates.
 */

/*
	This visitor performs most of the checks in our context analysis:
	-	Constructs the symbol table.
	-	Performs type checks according to the IC specifications.
	-	Enforces scope rules according to the IC specifications.
	-	Checks that the main method exists and is single.
	-	Checks for integer literal overflow (the 2^31 case that isn't checked in previous stages).
	-	Checks that break and continue are only used inside loops.
	-	Checks that the 'this' keyword is not used in a static method.
 */
public class ContextAnalysisVisitor implements Visitor {

    private TypeTable typeTable;
    private SymbolTable globalSymTable;
    private ContextError error;
    
    private Set<Expression> overflowLiterals;
    private boolean foundMain;
    private int nestedWhileCount;

    private Method mainMethod;

    public ContextAnalysisVisitor(TypeTable typeTable, String rootSymbolTableName) {
        this.error = null;
        this.typeTable = typeTable;
        this.globalSymTable = new SymbolTable(SymbolTableKind.Global, rootSymbolTableName, null, null);

        this.foundMain = false;
        this.overflowLiterals = new HashSet<Expression>();
        this.nestedWhileCount = 0;
    }

    public SymbolTable getGlobalSymTable() {
        return globalSymTable;
    }

    @Override
    public Object visit(Program program) {
        for (ICClass cls : program.getClasses()) {
            SymbolTable parentScope;
            if (cls.hasSuperClass()) {
                parentScope = cls.getSuperContextType().getIcClass().getClassSymbolTable();
            } else {
                parentScope = this.globalSymTable;
            }
            
            SymbolTable classSymbolTable = new SymbolTable(SymbolTableKind.Class, cls.getName(), parentScope, cls);
            Symbol clsSym = new Symbol(cls.getName(), cls.getContextType(), Kind.Class, cls);
            cls.setSymbolTableEntry(clsSym);
            this.globalSymTable.addSymbol(clsSym);

            cls.setEnclosingScope(this.globalSymTable);
            cls.setClassSymbolTable(classSymbolTable);

            //populate all fields and methods for use later inside methods
            for (Method method : cls.getMethods()) {
                MethodType mt = (MethodType)method.getContextType();
                if (classSymbolTable.search(method.getName(), false) != null) {
                    this.error = new ContextError("Identifier '" + method.getName() +
                                                    "' was already defined.", method.getLine());
                    return null;
                }
                
                Kind methodKind = Kind.StaticMethod;
                if (method instanceof VirtualMethod) {
                    methodKind = Kind.VirtualMethod;
                }

                // in case symbol is found, it is defined in the parentScope class.
                Symbol parentSymbol = classSymbolTable.search(method.getName(), true);
                if (parentSymbol != null && (parentSymbol.getKind() != methodKind || parentSymbol.getType() != mt)) {
                    this.error = new ContextError("Identifier '" + method.getName() +
                                                        "' was defined in the super class.", method.getLine());
                    return null;
                }

                // method is overridden
                if (parentSymbol != null && parentSymbol.getNode() instanceof Method) {
                    method.setOverriddenMethod((Method)parentSymbol.getNode());
                }

                Symbol methodSym = new Symbol(method.getName(), mt, methodKind, method);
                classSymbolTable.addSymbol(methodSym);
                method.setSymbolTableEntry(methodSym);
                this.checkIfMain(method);
                if (this.error != null) return null;
            }

            for (Field field : cls.getFields()) {
                if (classSymbolTable.search(field.getName(), true) != null) {
                    this.error = new ContextError("Identifier '" + field.getName() +
                            "' was already defined.", field.getLine());
                    return null;
                }

                Symbol fieldSym = new Symbol(field.getName(), field.getType().getContextType(), Kind.Field, field);
                field.setSymbolTableEntry(fieldSym);
                classSymbolTable.addSymbol(fieldSym);
            }

            //$this is used later for checking virtual method type tables.
            classSymbolTable.addSymbol(new Symbol("$this", cls.getContextType(), Kind.This, cls));
        }

        // recursive calls to classes
        for (ICClass cls : program.getClasses()) {
            cls.accept(this);
            if (this.error != null) return null;
        }

        // overflow check
        if (this.overflowLiterals.size() > 0) {
            Literal literal = ((Literal)this.overflowLiterals.iterator().next());
            this.error = new ContextError("Literal overflow in: " + literal.getValue() + '.', literal.getLine());
            return null;
        }

        if (this.foundMain) {
            program.setMain(this.mainMethod);
        } else {
            this.error = new ContextError("main function not found.", 0);
            return null;
        }

        return this.globalSymTable;
    }

    @Override
    public Object visit(ICClass icClass) {
        SymbolTable classSymTable = icClass.getClassSymbolTable();
        for (Method method : icClass.getMethods()) {
            method.setEnclosingScope(classSymTable);
            method.accept(this);
            if (this.error != null) return null;
        }

        for (Field field : icClass.getFields()) {
            field.setEnclosingScope(classSymTable);
            field.accept(this);
            if (this.error != null) return null;
        }

        return classSymTable;
    }

    @Override
    public Object visit(Field field) {
        // fields are added at program level.
        return null;
    }

    public Object visitMethod(Method method, boolean isStatic) {
        SymbolTable classSymbolTable = method.getEnclosingScope();
        SymbolTable methodSymTable = new SymbolTable(SymbolTableKind.Method, method.getName(), classSymbolTable, method);
        methodSymTable.setStatic(isStatic);

        method.setMethodSymTable(methodSymTable);

        //$ret is used later for type checking return statements.
        Symbol retSym = new Symbol("$ret", method.getType().getContextType(), Kind.ReturnVar, method);
        methodSymTable.addSymbol(retSym);
        method.setRetSymbolTableEntry(retSym);

        for (Formal formal : method.getFormals()) {
            formal.setEnclosingScope(methodSymTable);
            formal.accept(this);
            if (this.error != null) return null;
        }

        for (Statement statement : method.getStatements()) {
            statement.setEnclosingScope(methodSymTable);
            statement.accept(this);
            if (this.error != null) return null;
        }

        return methodSymTable;
    }

    @Override
    public Object visit(VirtualMethod method) {
        return visitMethod(method, false);
    }

    @Override
    public Object visit(StaticMethod method) {
        return visitMethod(method, true);
    }

    @Override
    public Object visit(LibraryMethod method) {
        return visitMethod(method, true);
    }

    @Override
    public Object visit(Formal formal) {
        Symbol formalSym = new Symbol(formal.getName(), formal.getType().getContextType(), Kind.Parameter, formal);
        formal.getEnclosingScope().addSymbol(formalSym);
        formal.setSymbolTableEntry(formalSym);
        return null;
    }

    @Override
    public Object visit(PrimitiveType type) {
        return null;  //To change body of implemented methods use File | Settings | File Templates.
    }

    @Override
    public Object visit(UserType type) {
        return null;  //To change body of implemented methods use File | Settings | File Templates.
    }

    @Override
    public Object visit(Assignment assignment) {
        Location var = assignment.getVariable();
        Expression val = assignment.getAssignment();

        var.setEnclosingScope(assignment.getEnclosingScope());
        val.setEnclosingScope(assignment.getEnclosingScope());

        var.accept(this);
        if (this.error != null) return null;
        val.accept(this);
        if (this.error != null) return null;
        
        IC.Types.Type varType = var.getContextType();
        IC.Types.Type valType = val.getContextType();

        if (!valType.subtypeOf(varType)) {
            this.error = new ContextError("Cannot assign '" + valType.getName() +
                                                "' into '" + varType.getName() + "'.", assignment.getLine());
            return null;
        }

        return null;
    }

    @Override
    public Object visit(CallStatement callStatement) {
        Call call = callStatement.getCall();
        call.setEnclosingScope(callStatement.getEnclosingScope());

        call.accept(this);

        return null;
    }

    @Override
    public Object visit(Return returnStatement) {
        Symbol retSym = returnStatement.getEnclosingScope().search("$ret", true);
        returnStatement.setSymbolTableEntry(retSym);
        IC.Types.Type methodReturnType = retSym.getType();
        if (returnStatement.hasValue()) {
            Expression val = returnStatement.getValue();
            val.setEnclosingScope(returnStatement.getEnclosingScope());
            val.accept(this);
            if (this.error != null) return null;

            if (!val.getContextType().subtypeOf(methodReturnType)) {
                this.error = new ContextError("Returning '" + val.getContextType().getName() +
                                                "' instead of '" + methodReturnType.getName() + "'.", returnStatement.getLine());
                return null;
            }

        } else if (methodReturnType != this.typeTable.getVoidType()) {
            this.error = new ContextError("Must return a value of type '" + methodReturnType.getName() +
                                                "'.", returnStatement.getLine());
            return null;
        }

        return null;
    }

    @Override
    public Object visit(If ifStatement) {
        Expression cond = ifStatement.getCondition();
        Statement op = ifStatement.getOperation();
        cond.setEnclosingScope(ifStatement.getEnclosingScope());
        op.setEnclosingScope(ifStatement.getEnclosingScope());

        cond.accept(this);
        if (this.error != null) return null;

        op.accept(this);
        if (this.error != null) return null;

        if (ifStatement.hasElse()) {
            Statement elseOp = ifStatement.getElseOperation();
            elseOp.setEnclosingScope(ifStatement.getEnclosingScope());

            elseOp.accept(this);
            if (this.error != null) return null;
        }

        if (cond.getContextType() != this.typeTable.getBoolType()) {
            this.error = new ContextError("If condition must be boolean.", cond.getLine());
            return null;
        }

        return null;
    }

    @Override
    public Object visit(While whileStatement) {
        Expression cond = whileStatement.getCondition();
        Statement op = whileStatement.getOperation();

        cond.setEnclosingScope(whileStatement.getEnclosingScope());
        op.setEnclosingScope(whileStatement.getEnclosingScope());

        cond.accept(this);
        if (this.error != null) return null;
        this.nestedWhileCount++;
        op.accept(this);
        this.nestedWhileCount--;
        if (this.error != null) return null;

        if (cond.getContextType() != this.typeTable.getBoolType()) {
            this.error = new ContextError("While condition must be boolean.", cond.getLine());
            return null;
        }

        return null;
    }

    @Override
    public Object visit(Break breakStatement) {
        if (this.nestedWhileCount == 0) {
            this.error = new ContextError("Break statement allowed only inside a loop.", breakStatement.getLine());
        }

        return null;
    }

    @Override
    public Object visit(Continue continueStatement) {
        if (this.nestedWhileCount == 0) {
            this.error = new ContextError("Continue statement allowed only inside a loop.", continueStatement.getLine());
        }

        return null;
    }

    @Override
    public Object visit(StatementsBlock statementsBlock) {
        SymbolTable parent = statementsBlock.getEnclosingScope();
        SymbolTable blockSymTable = new SymbolTable(SymbolTableKind.Block, "statement block in " + parent.getName(), parent, statementsBlock);

        for (Statement stm : statementsBlock.getStatements()) {
            stm.setEnclosingScope(blockSymTable);
            stm.accept(this);
            if (this.error != null) return null;
        }

        return null;
    }

    @Override
    public Object visit(LocalVariable localVariable) {
        // first check scope rules
        SymbolTable scope = localVariable.getEnclosingScope();
        Symbol sym = scope.search(localVariable.getName(), false);
        if (sym != null) {
            this.error = new ContextError("Variable '" + localVariable.getName() + "' is already defined.", localVariable.getLine());
            return null;
        }

        IC.Types.Type varType = localVariable.getType().getContextType();

        if (localVariable.hasInitValue()) {
            Expression val = localVariable.getInitValue();
            val.setEnclosingScope(scope);
            val.accept(this);
            if (this.error != null) return null;

            IC.Types.Type valType = val.getContextType();
            if (!valType.subtypeOf(varType)) {
                this.error = new ContextError("Cannot assign '" + valType.getName() +
                                                    "' into '" + varType.getName() + "'.", localVariable.getLine());
                return null;
            }
        }

        Symbol varSym = new Symbol(localVariable.getName(), varType, Kind.Var, localVariable);
        localVariable.setSymbolTableEntry(varSym);
        scope.addSymbol(varSym);

        return null;
    }

    private SymbolTable getClassScopeFromLocation(Expression location, SymbolTable enclosingScope) {

        location.setEnclosingScope(enclosingScope);
        location.accept(this);
        if (error != null) return null;

        IC.Types.Type expType = location.getContextType();
        if (!(expType instanceof ClassType)) {
            this.error = new ContextError("Cannot dereference the non-class type '" + expType.getName() + "'.", location.getLine());
            return null;
        }

        ClassType cLocType = (ClassType)expType;
        return cLocType.getIcClass().getClassSymbolTable();
    }

    @Override
    public Object visit(VariableLocation location) {
        SymbolTable scope = location.getEnclosingScope();

        if (location.isExternal()) {
            //if this is an external var (field), use the correct class's scope.
            scope = getClassScopeFromLocation(location.getLocation(), scope);
            if (this.error != null) return null;
        }

        Symbol varSymbol = scope.search(location.getName(), true);
        if (varSymbol == null || !varSymbol.getKind().isAllowedAsLocation()) {
            this.error = new ContextError("Unknown variable or field '" + location.getName() + "'.", location.getLine());
            return null;
        }

        location.setSymbolTableEntry(varSymbol);
        location.setContextType(varSymbol.getType());

        return null;
    }

    @Override
    public Object visit(ArrayLocation location) {
        Expression arr = location.getArray();
        Expression ind = location.getIndex();

        arr.setEnclosingScope(location.getEnclosingScope());
        arr.accept(this);
        if (this.error != null) return null;

        ind.setEnclosingScope(location.getEnclosingScope());
        ind.accept(this);
        if (this.error != null) return null;

        if (!(arr.getContextType() instanceof ArrayType)) {
            this.error = new ContextError("Cannot use the '[' operator on non array types.", location.getLine());
            return null;
        }

        if (ind.getContextType() != this.typeTable.getIntType()) {
            this.error = new ContextError("Array index must be an integer.", location.getLine());
            return null;
        }

        ArrayType arrType = (ArrayType)arr.getContextType();
        location.setContextType(arrType.getElementType());

        return null;
    }

    private void visitCall(SymbolTable scope, Symbol methodSym, Call call) {
        MethodType mt = (MethodType)methodSym.getType();
        IC.Types.Type[] parameterTypes = mt.getParameterTypes();
        List<Expression> args = call.getArguments();
        if (parameterTypes.length != args.size()) {
            this.error = new ContextError("Incorrect number of arguments " +
                    " when calling method '" + call.getName() + "'.", call.getLine());
            return;
        }

        for (int i = 0; i < parameterTypes.length; i++) {
            Expression argI = args.get(i);
            argI.setEnclosingScope(call.getEnclosingScope());
            argI.accept(this);
            if (this.error != null) return;

            if (!argI.getContextType().subtypeOf(parameterTypes[i])) {
                this.error = new ContextError("Incorrect parameter type given, parameter number: " + (i + 1) ,
                        call.getLine());
                return;
            }
        }

        call.setContextType(mt.getReturnType());
    }

    @Override
    public Object visit(StaticCall call) {
        SymbolTable scope = call.getEnclosingScope();
        SymbolTable classScope = call.getClassContextType().getIcClass().getClassSymbolTable();
        
        Symbol sym = classScope.search(call.getName(), true);
        if (sym == null || sym.getKind() != Kind.StaticMethod) {
            this.error = new ContextError("Unknown static method '" + call.getName() + "'.", call.getLine());
            return null;
        }

        call.setSymbolTableEntry(sym);
        visitCall(scope, sym, call);

        return null;
    }

    @Override
    public Object visit(VirtualCall call) {
        SymbolTable scope = call.getEnclosingScope();

        SymbolTable.SearchMode sm = SymbolTable.SearchMode.Default;
        if (call.isExternal()) {
            //if this is an external call, use the correct class's scope.
            scope = getClassScopeFromLocation(call.getLocation(), scope);
            if (this.error != null) return null;

            sm = SymbolTable.SearchMode.NonStaticOnly;
        }


        Symbol sym = scope.search(call.getName(), true, sm);
        if (sym == null || (sym.getKind() != Kind.StaticMethod && sym.getKind() != Kind.VirtualMethod)) {
            this.error = new ContextError("Unknown method '" + call.getName() + "'.", call.getLine());
            return null;
        }

        call.setSymbolTableEntry(sym);
        visitCall(scope, sym, call);

        return null;
    }

    @Override
    public Object visit(This thisExpression) {
        Symbol thisSym = thisExpression.getEnclosingScope().search("$this", true);

        if (thisSym == null) {
            this.error = new ContextError("The keyword 'this' can only be used inside virtual methods.", thisExpression.getLine());
            return null;
        }

        thisExpression.setSymbolTableEntry(thisSym);
        thisExpression.setContextType(thisSym.getType());
        return null;
    }

    @Override
    public Object visit(NewClass newClass) {
        //newClass already has ContextType, treated in BuildTypeTable.
        return null;
    }

    @Override
    public Object visit(NewArray newArray) {
        Expression size = newArray.getSize();
        size.setEnclosingScope(newArray.getEnclosingScope());
        size.accept(this);
        if (this.error != null) return null;

        if (size.getContextType() != this.typeTable.getIntType()) {
            this.error = new ContextError("Array size must be an integer.", newArray.getLine());
            return null;
        }

        //newArray already has ContextType, treated in BuildTypeTable.

        return null;
    }

    @Override
    public Object visit(Length length) {
        Expression arr = length.getArray();
        arr.setEnclosingScope(length.getEnclosingScope());
        arr.accept(this);
        if (this.error != null) return null;

        if (!(arr.getContextType() instanceof ArrayType)) {
            this.error = new ContextError("'length' keyword must be used with arrays.", length.getLine());
            return null;
        }

        length.setContextType(this.typeTable.getIntType());

        return null;
    }

    @Override
    public Object visit(MathBinaryOp binaryOp) {
        Expression fop = binaryOp.getFirstOperand();
        Expression sop = binaryOp.getSecondOperand();

        fop.setEnclosingScope(binaryOp.getEnclosingScope());
        fop.accept(this);
        if (this.error != null) return null;
        
        
        sop.setEnclosingScope(binaryOp.getEnclosingScope());
        sop.accept(this);
        if (this.error != null) return null;

        IC.Types.Type fopType = fop.getContextType();
        IC.Types.Type sopType = sop.getContextType();

        if (binaryOp.getOperator() == BinaryOps.PLUS &&
                fopType == this.typeTable.getStringType()
                &&
                sopType == this.typeTable.getStringType()) {
            binaryOp.setContextType(this.typeTable.getStringType());
        } else {
            if (fopType != this.typeTable.getIntType() ||
                    sopType != this.typeTable.getIntType()) {

                String operandTypes = "integers";
                if (binaryOp.getOperator() == BinaryOps.PLUS) {
                    operandTypes += " or strings";
                }
                
                this.error = new ContextError("when using the " + binaryOp.getOperator().getDescription() +
                        " operator, the operands must be " + operandTypes + "." , binaryOp.getLine());
            }

            binaryOp.setContextType(this.typeTable.getIntType());
        }

        return null;
    }

    @Override
    public Object visit(LogicalBinaryOp binaryOp) {

        Expression left = binaryOp.getFirstOperand();
        Expression right = binaryOp.getSecondOperand();

        right.setEnclosingScope(binaryOp.getEnclosingScope());
        left.setEnclosingScope(binaryOp.getEnclosingScope());

        left.accept(this);
        if (error != null) return null;

        right.accept(this);
        if (error != null) return null;

        IC.Types.Type leftType = left.getContextType();
        IC.Types.Type rightType = right.getContextType();

        BinaryOps op = binaryOp.getOperator();

        if (op == BinaryOps.LAND || op == BinaryOps.LOR) {
            if (leftType != this.typeTable.getBoolType() || rightType !=
                    this.typeTable.getBoolType()) {
                this.error = new ContextError("when using the " + op.getDescription() +
                        " operator, the operands must be booleans." , binaryOp.getLine());
                return null;
            }

        } else if (op == BinaryOps.GT || op == BinaryOps.GTE ||
                    op == BinaryOps.LT || op == BinaryOps.LTE) {
            if (leftType != this.typeTable.getIntType() || rightType != this.typeTable.getIntType()) {
                this.error = new ContextError("when using the " + op.getDescription() +
                        " operator, the operands must be integers." , binaryOp.getLine());
                return null;
            }
        } else if (op == BinaryOps.EQUAL || op == BinaryOps.NEQUAL) {
            if (!leftType.subtypeOf(rightType) && !rightType.subtypeOf(leftType)) {
                this.error = new ContextError("when using the " + op.getDescription() +
                        " operator, the operands must have inheritance relationship." , binaryOp.getLine());
                return null;
            }
        }
        
        binaryOp.setContextType(this.typeTable.getBoolType());
        return null;
    }

    @Override
    public Object visit(MathUnaryOp unaryOp) {
        Expression operand =  unaryOp.getOperand();
        operand.setEnclosingScope(unaryOp.getEnclosingScope());

        operand.accept(this);
        if (this.error != null) return null;
        
        if (operand.getContextType() != this.typeTable.getIntType()) {
            this.error = new ContextError("when using the " + unaryOp.getOperator().getDescription() +
                    " operator, the operand must be an integer." , unaryOp.getLine());
            return null;
        }

        if (this.overflowLiterals.contains(operand)) {
            // not overflow -2^31
            this.overflowLiterals.remove(operand);
        }

        unaryOp.setContextType(this.typeTable.getIntType());
        return null;
    }

    @Override
    public Object visit(LogicalUnaryOp unaryOp) {
        Expression operand =  unaryOp.getOperand();
        operand.setEnclosingScope(unaryOp.getEnclosingScope());

        operand.accept(this);
        if (this.error != null) return null;

        if (operand.getContextType() != this.typeTable.getBoolType()) {
            this.error = new ContextError("when using the " + unaryOp.getOperator().getDescription() +
                    " operator, the operand must be a boolean." , unaryOp.getLine());
            return null;
        }

        unaryOp.setContextType(this.typeTable.getBoolType());
        return null;
    }

    @Override
    public Object visit(Literal literal) {
        switch (literal.getType()) {
            case INTEGER:
                literal.setContextType(this.typeTable.getIntType());

                // check for overflow
                if (((Long)literal.getValue()) == 0x80000000L) {
                    this.overflowLiterals.add(literal);
                }
                break;
            case STRING:
                literal.setContextType(this.typeTable.getStringType());
                break;
            case TRUE:
                literal.setContextType(this.typeTable.getBoolType());
                break;
            case FALSE:
                literal.setContextType(this.typeTable.getBoolType());
                break;
            case NULL:
                literal.setContextType(this.typeTable.getNullType());
                break;
        }

        return null;
    }

    @Override
    public Object visit(ExpressionBlock expressionBlock) {
        Expression exp = expressionBlock.getExpression();
        exp.setEnclosingScope(expressionBlock.getEnclosingScope());

        exp.accept(this);
        if (this.error != null) return null;

        expressionBlock.setContextType(exp.getContextType());
        return null;
    }

    public ContextError getError() {
        return error;
    }

    private void checkIfMain(Method method) {
        MethodType methodType = method.getContextType();

        boolean isMain = (method.getName().equals("main") &&
                methodType.getParameterTypes().length == 1 &&
                methodType.getParameterTypes()[0] == this.typeTable.getArrayType(this.typeTable.getStringType()) &&
                methodType.getReturnType() == this.typeTable.getVoidType() &&
                method instanceof StaticMethod);
        if (this.foundMain && isMain) {
            this.error = new ContextError("Main method defined more than once.", method.getLine());
        }

        if (isMain) {
            this.foundMain = isMain;
            this.mainMethod = method;
        }
    }
}
